package com.agent.platform.client;

import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Mono;

import java.time.Duration;
import java.util.List;
import java.util.Map;

/**
 * 上下文管理器客户端 - 对应Node.js版本的ContextManager接口
 */
@Slf4j
@Component
public class ContextManagerClient {

    private final WebClient webClient;
    private final ObjectMapper objectMapper;

    public ContextManagerClient(WebClient.Builder webClientBuilder, ObjectMapper objectMapper) {
        this.webClient = webClientBuilder.build();
        this.objectMapper = objectMapper;
    }

    /**
     * 获取会话消息
     */
    public Mono<GetMessagesResponse> getMessages(
        String contextManagerUrl,
        String sessionId,
        String agentId,
        int limit,
        int offset
    ) {
        GetMessagesRequest request = new GetMessagesRequest();
        request.setSessionId(sessionId);
        request.setAgentId(agentId);
        request.setLimit(limit);
        request.setOffset(offset);

        return webClient.post()
            .uri(contextManagerUrl + "/messages/get")
            .bodyValue(request)
            .retrieve()
            .bodyToMono(GetMessagesResponse.class)
            .timeout(Duration.ofSeconds(10))
            .doOnError(e -> log.error("Failed to get messages from context manager", e));
    }

    /**
     * 添加消息
     */
    public Mono<Void> addMessage(
        String contextManagerUrl,
        String sessionId,
        String agentId,
        MessageData message
    ) {
        AddMessageRequest request = new AddMessageRequest();
        request.setSessionId(sessionId);
        request.setAgentId(agentId);
        request.setMessage(message);

        return webClient.post()
            .uri(contextManagerUrl + "/messages/add")
            .bodyValue(request)
            .retrieve()
            .bodyToMono(Void.class)
            .timeout(Duration.ofSeconds(10))
            .doOnError(e -> log.error("Failed to add message to context manager", e));
    }

    /**
     * 健康检查
     */
    public Mono<Boolean> healthCheck(String contextManagerUrl) {
        return webClient.get()
            .uri(contextManagerUrl + "/health")
            .retrieve()
            .bodyToMono(String.class)
            .map(response -> true)
            .timeout(Duration.ofSeconds(5))
            .onErrorReturn(false);
    }

    @Data
    public static class GetMessagesRequest {
        private String sessionId;
        private String agentId;
        private int limit;
        private int offset;
    }

    @Data
    public static class GetMessagesResponse {
        private String sessionId;
        private List<MessageData> messages;
        private int totalCount;
        private boolean hasMore;
    }

    @Data
    public static class AddMessageRequest {
        private String sessionId;
        private String agentId;
        private MessageData message;
    }

    @Data
    public static class MessageData {
        private String id;
        private String role; // user, assistant, system, tool
        private String content;
        private Map<String, Object> metadata;
        private String timestamp;
    }
}
